import matplotlib.pyplot as plt #数据可视化
import jieba.posseg as posseg
import wordcloud #分词
from wordcloud import WordCloud,ImageColorGenerator,STOPWORDS #词云，颜色生成器，停止词
import numpy as np #科学计算
from PIL import Image #处理图片
import csv, os, collections
from datetime import datetime
s = '2021-08-07'


def read_file(start_date, end_date=datetime.now()):
    res = ''
    for lastname in os.listdir(s):
        filename = os.path.join(s, lastname)
        # print(filename)
        with open(filename, newline='', encoding='utf-8') as csvfile:
            reader = csv.DictReader(csvfile)
            for row in reader:
                if row['quote'] == '' or row['is_rt'] == 'True':
                    continue
                post_date = datetime.strptime(row['datetime'], "%Y年%m月%d日 %H:%M")
                if post_date < start_date or post_date > end_date:
                    continue
                res += row['quote'] + '\n\n'
    # print(res)
    return res


def read_stopwords(files):
    stopwords = []
    for f in files:
        with open(f, encoding='utf-8') as fp:
            stopword = fp.readlines()
        stopwords += [s.replace('\n', '') for s in stopword]
    return set(stopwords)


def word_cloud(text):
    stopwords=read_stopwords(['cn_stopwords.txt', 'hit_stopwords.txt', 'user_stopwords.txt'])
    
    wordlist = posseg.cut(text)#切割词语
    wl = [word for word in wordlist if word.word not in stopwords]
    wl = [word for word in wl if word.flag.startswith('n')]
    ol = [w.word for w in wl]
    word_counts = collections.Counter(ol)
    backgroud = np.array(Image.open('test.png')) 

    stopwords = []
    wc = WordCloud(width=1800, height=1600,
        background_color='white',
        mode='RGB', 
        mask=backgroud, #添加蒙版，生成指定形状的词云，并且词云图的颜色可从蒙版里提取
        max_words=700,
        stopwords=STOPWORDS,
        font_path='C:\WINDOWS\FONTS\MSYH.ttc',
        max_font_size=150,
        relative_scaling=0.6, #设置字体大小与词频的关联程度为0.4
        random_state=50, 
        scale=2 
    ).generate_from_frequencies(word_counts)

    image_color = ImageColorGenerator(backgroud)#设置生成词云的颜色，如去掉这两行则字体为默认颜色
    wc.recolor(color_func=image_color)

    plt.imshow(wc) #显示词云
    plt.axis('off') #关闭x,y轴
    plt.show()#显示
    wc.to_file('qqzone_wordcloud.jpg') #保存词云图


def main():
    start_text = '2013-01-01'
    start_date = datetime.strptime(start_text, '%Y-%m-%d')
    end_text = '2018-01-01'
    end_date = datetime.strptime(end_text, '%Y-%m-%d')
    text = read_file(start_date)
    word_cloud(text)
 
if __name__ == '__main__':
    main()
